import os
import cv2
import xml.etree.ElementTree as ET


def get_coordinate(anno_path):
    coor_list = []
    with open(anno_path, "r", encoding='UTF-8') as xml_file:
        tree = ET.parse(xml_file)
        root = tree.getroot()
        # print(root.tag)
        # print(root.attrib)
        for obj in root.iter('object'):
            tem_coordinate = []
            xml_box = obj.find('bndbox')
            tem_coordinate.append(int(xml_box.find('xmin').text))
            tem_coordinate.append(int(xml_box.find('ymin').text))
            tem_coordinate.append(int(xml_box.find('xmax').text))
            tem_coordinate.append(int(xml_box.find('ymax').text))
            coor_list.append(tem_coordinate)
    return coor_list


def crop_image(source_file, dest_dir, coor_list):
    picture = cv2.imread(source_file)
    # pic_config = picture.shape
    # height = pic_config[0]
    # width = pic_config[1]
    # color_channel = pic_config[2]
    _, full_name = os.path.split(source_file)
    file_prefix, _ = os.path.splitext(full_name)
    count = 0
    for tem_coor in coor_list:
        count += 1
        tem_dest_file = os.path.join(dest_dir, file_prefix + '_' + str(count) + '.jpg')
        print("processing ", tem_dest_file)
        xmin = tem_coor[0]
        ymin = tem_coor[1]
        xmax = tem_coor[2]
        ymax = tem_coor[3]
        crop_img = picture[ymin:ymax, xmin:xmax]
        cv2.imwrite(tem_dest_file, crop_img)


if __name__ == '__main__':
    root_dir = "/home/kaijia/algo-env/datasets"
    dest_dir = "/home/kaijia/algo-env/Utils/qr_code/img_qr"
    if not os.path.exists(dest_dir):
        os.mkdir(dest_dir)
    anno_dir = os.path.join(root_dir, 'Annotations')
    img_dir = os.path.join(root_dir, 'Images')
    anno_list = os.listdir(anno_dir)
    for anno in anno_list:
        file_prefix = anno[:-4]
        img_path = os.path.join(img_dir, file_prefix + '.jpg')
        anno_path = os.path.join(anno_dir, anno)
        coor_list = get_coordinate(anno_path)
        crop_image(img_path, dest_dir, coor_list)
